import os
import sys

import loralib as lora
from wenet.lora.lora_utils import lora_sub

import logging

def set_lora(model, config):
    if 'encoder_lora' in config.keys():
        lora_conf = config['encoder_lora']
        lora_keys = lora_conf['lora_keys']
        if len(lora_keys) != 0:
            logging.info(f'Encoder uses lora! key is: {lora_keys}')
            model.encoder = lora_sub(model.encoder, lora_keys, 
                                     lora_r=lora_conf['r'], 
                                     lora_alpha=lora_conf['alpha'],
                                     lora_dropout=lora_conf['dropout'], 
                                     merge_weights=False, 
                                     use_adalora=lora_conf['use_adalora'])
            #model.ctc = lora_sub(model.ctc, lora_keys, 
            #                     lora_r=lora_conf['r'], 
            #                     lora_alpha=lora_conf['alpha'],
            #                     lora_dropout=lora_conf['dropout'], 
            #                     merge_weights=False, 
            #                     use_adalora=lora_conf['use_adalora'])
        else:
            logging.info(f'Encoder is freeze!')
        lora.mark_only_lora_as_trainable(model.encoder)
        #lora.mark_only_lora_as_trainable(model.ctc)
    else:
        logging.info(f'Encoder will finetune all params!')

    if 'llm_lora' in config.keys():
        lora_conf = config['llm_lora']
        lora_keys = lora_conf['lora_keys']
        if len(lora_keys) != 0:
            logging.info(f'LLM uses lora! key is: {lora_keys}')
            model.llm = lora_sub(model.llm, lora_keys, 
                                 lora_r=lora_conf['r'], 
                                 lora_alpha=lora_conf['alpha'],
                                 lora_dropout=lora_conf['dropout'], 
                                 merge_weights=False, 
                                 use_adalora=lora_conf['use_adalora'])
        else:
            logging.info(f'LLM is freeze! This is a mistake!!!')
        lora.mark_only_lora_as_trainable(model.llm)
    else:
        logging.info(f'LLM will finetune all params!')
